/*
 *  Copyright 2013 National Institute of Advanced Industrial Science and Technology
 *  
 *  Licensed under the Apache License, Version 2.0 (the "License");
 *  you may not use this file except in compliance with the License.
 *  You may obtain a copy of the License at
 *  
 *      http://www.apache.org/licenses/LICENSE-2.0
 *  
 *  Unless required by applicable law or agreed to in writing, software
 *  distributed under the License is distributed on an "AS IS" BASIS,
 *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *  See the License for the specific language governing permissions and
 *  limitations under the License.
 */

package org.sss.common;

import java.util.Arrays;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;

public class WriterCounter {
  private class Counter {
    public Counter() {
      this.count  = new long[nOutputs];
      this.nBytes = new long[nOutputs];
      Arrays.fill(count, 0L);
      Arrays.fill(nBytes, 0L);

      map.put(Thread.currentThread(), this);
    }
    public final long[] count;
    public final long[] nBytes;
  }

  private int nOutputs;
  private ConcurrentMap<Thread, Counter> map;
  private ThreadLocal<Counter> counter;

  public WriterCounter(int n) {
    this.nOutputs = n;
    this.map = new ConcurrentHashMap<Thread, Counter>();
    this.counter = new ThreadLocal<Counter>() {
      @Override
      protected Counter initialValue() {
        return new Counter();
      }
    };
  }

  public void countup(int n, long count, long nBytes) {
    Counter c = counter.get();
    c.count[n]  += count;
    c.nBytes[n] += nBytes;
  }

  public long[] getCount() {
    long[] cnt = new long[nOutputs];
    Arrays.fill(cnt, 0L);

    for (Counter c: map.values()) {
      for (int i = 0;i < nOutputs;++i) {
        cnt[i] += c.count[i];
      }
    }
    return cnt;
  }

  public long[] getNBytes() {
    long[] b = new long[nOutputs];
    Arrays.fill(b, 0L);

    for (Counter c: map.values()) {
      for (int i = 0;i < nOutputs;++i) {
        b[i] += c.nBytes[i];
      }
    }
    return b;
  }
}
